from python_ai.category.data.mnist_dataset import x_train, x_test, x_val, y_train, y_test, y_val, shape_, y_train_oh, y_test_oh, y_val_oh
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt
import os
import time
import io
from PIL import Image
from mlp_impl import MLP
import mlp_impl_sigmoid as mlps
import python_ai.category.tvts.tvts as tvts

# Hyper params
NAME = 'mnist_multi_clf_by_numpynet_mlp_softmax_ok_tvts_cli'
RSEED = 1
if RSEED is not None:
    np.random.seed(RSEED)
VER = 'v1.0'
N_BATCH_SIZE = 256
LR = 0.1  # learning rate (It should be higher for all sigmoid layers than all relu ones)
N_EPOCH = 100  # Please to be greater than 650 for a good model; to be 300 for a normal P-R curve
SAVE_FREQ = 10
PARENT_ID = 0
PARENT_EPOCH = 100
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)
os.makedirs(SAVE_DIR, exist_ok=True)
PID = os.getpid()
PYTHON = '/home/asuspei/miniconda3/envs/p37tf24/bin/python3'
TVTS_CLI = os.path.join(BASE_DIR, '..', '..', 'tvts', 'tvts_cli.py')
CLI_PREFIX = f'"{PYTHON}" "{TVTS_CLI}" "{NAME}" "{PID}" '


def arg_list_to_cli_str(xlist):
    xlist_q = [f'"{el}"' for el in xlist]
    cli_str = " ".join(xlist_q)
    return cli_str


# tvts_cli init
T_FLOAT = 'float'
T_INT = 'int'
T_STR = 'str'
x_init_params = [
    'ver', T_STR, VER,
    'batch_size', T_INT, N_BATCH_SIZE,
    'lr', T_FLOAT, LR,
    'n_epoch', T_INT, N_EPOCH,
    'save_freq', T_INT, SAVE_FREQ,
    'parent_id', T_INT, PARENT_ID,
    'parent_epoch', T_INT, PARENT_EPOCH,
]
args_cli = arg_list_to_cli_str(x_init_params)
cli = f'{CLI_PREFIX} init {args_cli}'
ret = os.system(cli)
if ret:
    raise Exception(f'TVTS CLI FAILED WITH RETURN VALUE {ret}: {cli}')
CLI_PATH_PREFIX = f'{NAME}_{PID}'
CKPT_PATH_PATH = os.path.join(BASE_DIR, '..', '..', 'tvts', '_save', 'tvts_cli.py', 'v1.0', f'{CLI_PATH_PREFIX}.txt')
if PARENT_ID == 0:
    CKPT_PATH = None
else:
    while not os.path.exists(CKPT_PATH_PATH):
        print('Waiting ...')
        time.sleep(0.2)
    with open(CKPT_PATH_PATH, 'r', encoding='utf8') as f:
        CKPT_PATH = f.read()
IS_TRAIN = 1


def get_saved_path_prefix(xepoch):
    return os.path.join(SAVE_DIR, CLI_PATH_PREFIX)


def get_saved_path(xepoch):
    prefix = get_saved_path_prefix(xepoch)
    return f'{prefix}.npy'


def get_curve_path(xsave_path):
    return f'{xsave_path}.png'


def run_on_data(model, label, x_data, y_data, n_batch_size, is_train, epoch):
    n = len(x_data)
    n_batch = int(np.ceil(n / n_batch_size))
    cost_avg = 0.
    metrics_avg = None
    pred = np.array([])
    for i, (x_batch, y_batch) in enumerate(mlps.dl(x_data, y_data, n_batch_size, is_train)):
        batch = i + 1
        h = model(x_batch)
        pred = np.append(pred, h)
        cost = model.cost(h, y_batch)
        metrics = model.metrics(y_batch, h)
        print(label, 'epoch:', epoch, 'batch:', batch, 'cost:', cost, 'metrics:', metrics)
        acc = metrics.tolist()[0][0]
        if is_train:
            # tvts_cli save_batch
            xi_params = x_init_params.copy()
            xi_params.extend([
                'cost', T_FLOAT, cost,
                'acc', T_FLOAT, acc,
                'epoch', T_INT, epoch,
                'batch', T_INT, batch,
            ])
            args_cli = arg_list_to_cli_str(xi_params)
            cli = f'{CLI_PREFIX} save_batch {args_cli}'
            ret = os.system(cli)
            if ret:
                raise Exception(f'TVTS CLI FAILED WITH RETURN VALUE {ret}: {cli}')
        cost_avg += cost / n_batch
        if metrics_avg is None:
            metrics_avg = np.zeros_like(metrics, dtype=np.float32)
        metrics_avg += metrics / n_batch
        if is_train:
            model.backward(y_batch)
            model.step()
    print(label, 'epoch:', epoch, 'cost_avg:', cost_avg, 'metrics:', metrics_avg)
    return cost_avg, metrics_avg, pred.reshape(*y_data.shape)


# Build model
sep('Model')
model = MLP(x_test.shape[1], [(128, 'sigmoid'), (64, 'sigmoid'), (32, 'sigmoid'), (10, 'linear')], LR)

# Train or load
if CKPT_PATH is not None:
    sep('Load weight')
    if not os.path.exists(CKPT_PATH):
        raise Exception(f'Check point file {CKPT_PATH} is not existed')

    print(f'Loading from {CKPT_PATH} ...')
    model.load_weights(CKPT_PATH)
    curve_save_path = get_curve_path(CKPT_PATH)
    print('Loaded.')

if IS_TRAIN:
    sep('Train')
    cost_his, metrics_his = [], []
    cost_his_val, metrics_his_val = [], []
    for i in range(N_EPOCH):
        epoch = i + 1
        cost_avg, metrics_avg, *_ = run_on_data(model, 'train', x_test, y_test_oh, N_BATCH_SIZE, True, epoch)
        metrics_avg = metrics_avg.tolist()
        cost_his.append(cost_avg)
        metrics_his.append(metrics_avg)
        cost_avg_val, metrics_avg_val, *_ = run_on_data(model, 'val', x_test, y_test_oh, N_BATCH_SIZE, False, epoch)
        metrics_avg_val = metrics_avg_val.tolist()
        cost_his_val.append(cost_avg_val)
        metrics_his_val.append(metrics_avg_val)
        save_path = None
        if epoch % SAVE_FREQ == 0:
            save_prefix = get_saved_path_prefix(epoch)
            save_path = get_saved_path(epoch)
            print(f'Saving to {save_path} ...')
            model.save_weights(save_prefix)
            print('Saved.')

        # tvts_cli save_epoch
        xi_params = x_init_params.copy()
        xi_params.extend([
            'cost', T_FLOAT, cost_avg,
            'acc', T_FLOAT, metrics_avg[0][0],
            'cost_val', T_FLOAT, cost_avg_val,
            'acc_val', T_FLOAT, metrics_avg_val[0][0],
            'epoch', T_INT, epoch
        ])
        if save_path is not None:
            xi_params.extend([
                'save_path', T_STR, save_path,
            ])
        args_cli = arg_list_to_cli_str(xi_params)
        cli = f'{CLI_PREFIX} save_epoch {args_cli}'
        ret = os.system(cli)
        if ret:
            raise Exception(f'TVTS CLI FAILED WITH RETURN VALUE {ret}: {cli}')
    metrics_his = np.array(metrics_his)
    metrics_his_val = np.array(metrics_his_val)

    # save training curve
    plt.figure(figsize=[12, 6])
    spr = 1
    spc = 2
    spn = 0

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(f'Cost(lr={LR})')
    plt.plot(cost_his, label='train')
    plt.plot(cost_his_val, label='val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Metrics')
    plt.plot(metrics_his[:, 0], label='acc_train')
    plt.plot(metrics_his_val[:, 0], label='acc_val')
    plt.legend()
    plt.grid()

    save_path = get_saved_path(epoch)
    curve_save_path = get_curve_path(save_path)
    plt.savefig(curve_save_path)
    plt.close()

# do test
sep('Test')
cost_avg, metrics_avg, pred = run_on_data(model, 'test', x_test, y_test_oh, N_BATCH_SIZE, False, 0)

# Draw testing demo in memory buf
plt.figure(figsize=[6, 6])
spr = 5
spc = 5
spn = 0
m = spr * spc
for i in range(m):
    spn += 1
    plt.subplot(spr, spc, spn)
    y = y_test[i, 0]
    img = x_test[i].reshape(*shape_)
    pr = pred[i].argmax()
    plt.axis('off')
    plt.title(f'{y}=>{pr}', color='black' if pr == y else 'red')
    plt.imshow(img)
buf = io.BytesIO()
buf.seek(0)
plt.savefig(buf, format='png')
img_demo = Image.open(buf)
plt.close()

# Show curve of cost and metrics in training if it exits and show the demo of test
if not IS_TRAIN and os.path.exists(curve_save_path):
    plt.figure(figsize=[15, 5])
    plt.axis('off')
    img2 = Image.open(curve_save_path)
    images = (img_demo, img2)
    widths, heights = zip(*(i.size for i in images))
    total_width = sum(widths)
    max_height = max(heights)
    new_im = Image.new('RGB', (total_width, max_height))
    x_offset = 0
    for im in images:
        new_im.paste(im, (x_offset, 0))
        x_offset += im.size[0]
    plt.imshow(new_im)
    plt.show()
elif IS_TRAIN:
    plt.figure(figsize=[15, 5])
    plt.axis('off')
    spr = 1
    spc = 3
    spn = 0

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(f'Cost(lr={LR})')
    plt.plot(cost_his, label='train')
    plt.plot(cost_his_val, label='val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('Metrics')
    plt.plot(metrics_his[:, 0], label='acc_train')
    plt.plot(metrics_his_val[:, 0], label='acc_val')
    plt.legend()
    plt.grid()

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.axis('off')
    plt.imshow(img_demo)
    buf.close()

    plt.show()
else:
    plt.figure(figsize=[6, 6])
    plt.axis('off')
    plt.imshow(img_demo)
    buf.close()
    plt.show()
